feat : a general probabilistic forecasting interface#700
feat : a general probabilistic forecasting interface#700Sir-Sloth-The-Lazy wants to merge 17 commits into
Conversation
Add abstract Forecaster.compute_training_loss returning a finished (loss, loss_components) pair, so each forecaster owns its complete training objective. ForecasterModule.training_step now only injects the configured scoring rule, interior mask and per_var_std, and logs the result. The deterministic ARForecaster loss is unchanged in value. Add the abstract ProbabilisticForecaster (sample_ensemble capability), ProbabilisticARForecaster (sequential sampled rollouts, trains on the configured score of the ensemble mean) and a minimal ProbabilisticForecasterModule whose validation samples an ensemble and logs the RMSE of the ensemble mean. Interface design from mllam#685.
|
@joeloskarsson @observingClouds , this is the implementation of the dicussion on issue #685 hope this is what you wanted ! 👀. |
observingClouds
left a comment
There was a problem hiding this comment.
Hi @Sir-Sloth-The-Lazy thanks for drafting this. I had a look at the discussion around the proposal in #685 and this PR. It looks well aligned. So far, I only have a few minor comments.
| trajectory. This class adds ensemble forecasting on top: unrolling | ||
| several trajectories and stacking them along an ensemble dimension. | ||
| The default training objective scores the ensemble mean with the | ||
| injected scoring rule; forecasters with model-specific objectives |
There was a problem hiding this comment.
"injected scoring rule". Can this be written more explicitly? So that it is clear to users where the scoring rule is set.
| states at each predicted step, used both as the prediction | ||
| targets and to overwrite boundary nodes during the rollouts. | ||
| Dims: same as one ensemble member. | ||
| score_fn : Callable |
There was a problem hiding this comment.
fn always reminds me as an abbreviation for filename. Maybe use score_func or score_metric?
There was a problem hiding this comment.
Or even just score.
There was a problem hiding this comment.
renamed to score_metric as per the suggestion.
|
@observingClouds , Hope the latest commit solves the concerns 😁 |
joeloskarsson
left a comment
There was a problem hiding this comment.
This looks quite nice I think, and is still pretty easy to follow. I had some outstanding thoughts regarding the connection to the deterministic version + some small things.
| init_states, | ||
| forcing_features, | ||
| target_states, | ||
| score_metric=self.loss, |
There was a problem hiding this comment.
If the loss calculation is now entirely handled by the Forecaster, I think it could make more sense that self.loss is instantiated and sits on the Forecaster, rather than being passed to it? What do you think, is there any problem with that?
There was a problem hiding this comment.
Good question, I don't think there's a hard problem with it, but it's a bit bigger than a one-liner so wanted to lay out the scope before doing it.
Right now self.loss isn't only feeding compute_training_loss ForecasterModule also calls it directly in validation_step/test_step (and for the spatial loss map) to report val/test loss, independent of whatever the forecaster's training objective ends up being. So "loss calculation entirely handled by the Forecaster" isn't quite true yet for the deterministic path the Module still owns a second, separate use of the same metric for reporting.
It also isn't uniform across the two modules we have: ProbabilisticForecasterModule.validation_step doesn't touch self.loss at all, it scores the ensemble mean with raw metrics.mse, and test_step isn't implemented there yet.
If we move it, the plan would be:
loss: str becomes a ctor arg on Forecaster (or ARForecaster), instantiated there as self.loss = metrics.get_metric(loss)
compute_training_loss drops the score_metric param entirely and reads self.loss internally
ForecasterModule.validation_step/test_step read self.forecaster.loss instead of self.loss
loss gets threaded through the forecaster constructors in train_model.py (currently only passed to ForecasterModule) and the ~15 test call sites that build ARForecaster/ProbabilisticARForecaster directly
None of that is a blocker, just confirming that's the scope before I make the change and whether ForecasterModule should keep a loss kwarg at all (for backward-compat / overriding the reporting metric independently of the training metric) or drop it entirely in favor of always reading it off the forecaster.
There was a problem hiding this comment.
I am not a fan of having the ForecasterModule get the loss as self.forecaster.loss and then using that. A correct division of responsibilities would mean that the ForecasterModule does not know how to compute the loss for a batch, and that should be delegated to the Forecaster. So instead of getting the loss and using it, I think any loss computation that is needed in the validation/test step should be delegated to the Forecaster. If that means that we need additional functionality in the Forecaster (e.g. other types of loss reduction, not all dims) than we should implement that there.
| target_states, | ||
| score_metric=self.loss, | ||
| interior_mask_bool=self.interior_mask_bool, | ||
| per_var_std=self.per_var_std, |
There was a problem hiding this comment.
Similarly for per_var_std, the main use case of this seems to be computing the loss, so would it also belong in the Forecaster then, that handles the loss?
There was a problem hiding this comment.
-
On score_metric/self.loss: Moved. ARForecaster (and ProbabilisticARForecaster through it) now takes aloss: str = "wmse"constructor argument and buildsself.loss = metrics.get_metric(loss)itself.compute_training_lossno longer takes ascore_metricparam, it readsself.lossinternally.ForecasterModule.training_stepnow just injects the interior mask; validation_step/test_step readself.forecaster.lossfor their reporting metric instead of owning a separate copy. -
On
per_var_std: Moved too, same reasoning. ARForecaster also takes an optionalconfig: NeuralLAMConfig = Noneconstructor argument and computesself.per_var_stdfrom it (same formula as beforediff_std / sqrt(feature_weights)), only when notself.predicts_std. Made config optional rather than required: since per_var_std is only needed for scoring, a forecaster built purely for inference (or the several existing tests that only exercise forward()) doesn't need to supply aNeuralLAMConfigat all it staysNoneand is simply never used, same as today'spredicts_std=Truecase.
There was a problem hiding this comment.
This sounds good, but what happens if we use StepPredictor that does not output a pred_std, but also no NeuralLAMConfig is given as an argument ( = None)? That seems to potentially cause issues.
I am not certain if we need to allow the config=None case, there should always be such a config present so we could require it.
There was a problem hiding this comment.
I see you noticed this as well, and fixed in 511a6d5. However, I don't think this check sits in the right spot. It checks the Forecaster configuration, so it should be in the Forecaster rather than the ForecastModule probably.
|
Thank you for taking out the time to review this @joeloskarsson means a lot. I will start working on the review now 😃 |
score_metric/per_var_std were injected into compute_training_loss by ForecasterModule and also used directly for val/test loss reporting, duplicating config the forecaster already needs for its own objective. ARForecaster/ProbabilisticARForecaster now own self.loss and self.per_var_std (computed from an optional config ctor arg), and ForecasterModule reads them off self.forecaster instead. Also trims the CHANGELOG entry for mllam#685 down to one sentence per review feedback.
A Forecaster built without config now silently has per_var_std=None when its predictor doesn't output its own std. Previously per_var_std was always computed by ForecasterModule itself, so this gap didn't exist; now that construction is split across two calls, catch it at ForecasterModule init instead of crashing at the first val/test step.
Each member's predicted std is its own, not a spread computed across the ensemble, so ensemble_std was a misleading name. Document on ProbabilisticForecaster that a per-member std makes the predictive distribution a mixture of Gaussians, and note in ProbabilisticARForecaster.compute_training_loss that averaging the per-member stds is a simplification of the true mixture variance (which also includes the spread between member means).
Co-authored-by: Joel Oskarsson <joel.oskarsson@outlook.com>
Scoring the ensemble mean with a pointwise metric only rewards the mean being right, giving the model no incentive to keep a calibrated spread, and risks training it to collapse to a point estimate. Redeclare compute_training_loss as abstract on ProbabilisticARForecaster instead of providing that as a default (it would otherwise silently fall back to ARForecaster's single-rollout objective via MRO, not even the ensemble mean). Concrete subclasses must define their own objective. Tests that only need an instantiable forecaster now use a local ConcreteProbabilisticARForecaster example (ensemble-mean scoring, moved out of the library code); a new test locks in that the base class itself cannot be instantiated.
Drop ProbabilisticARForecaster's ensemble_size constructor arg and the implicit num_members=None -> self.ensemble_size fallback in sample_ensemble; num_members is now always required. Baking a default member count into the forecaster's state was unnecessary now that compute_training_loss is abstract too (nothing in the shared base class path used it) and just adds an implicit default callers could silently rely on instead of deciding explicitly. The num_members < 1 validation moves from __init__ to sample_ensemble accordingly. ProbabilisticForecasterModule.eval_ensemble_size follows suit: it no longer defaults to None with a forecaster fallback, it's required. Test-only ConcreteProbabilisticARForecaster (used wherever a concrete probabilistic forecaster is needed for testing) gains its own train_num_members for the training objective, since deciding how many members to sample during training is now the concrete subclass's call.
Mirrors validation_step: samples eval_ensemble_size members and scores the ensemble mean, same as validation. Factored the shared sampling + scoring + logging into _ensemble_step(batch, phase) rather than duplicating the block, since validation_step and test_step differ only in their log-key prefix and which metrics dict collects the result. Overrides on_test_epoch_end (rather than inheriting ForecasterModule's) since this module's test_step doesn't populate spatial_loss_maps or plot examples - the inherited version would crash on torch.cat of an empty list.
| # per_var_std is normally computed from config; override directly since | ||
| # this test only cares about the loss computation, not standardization. |
There was a problem hiding this comment.
I'm not much of a fan of these inline comments and think the code should be clear enough to explain itself.
There was a problem hiding this comment.
Fair point, removed the assignment is clear on its own. Committed as 45ffdeb. Is there something more you want out of this ? I would really like to see, if this is enough or I can do more to make the code more readable.
| from .module import ForecasterModule | ||
|
|
||
|
|
||
| class ProbabilisticForecasterModule(ForecasterModule): |
There was a problem hiding this comment.
Instead of using inheritance here I would make an abstract base class, e.g. BaseForecasterModule, which implements expected methods as abstractmethods, and then renaming the ForecasterModule as DeterministicForecasterModule. And then ProbabilisticForecasterModule and DeterministicForecasterModule become an implementation of this base class.
Doing it this way will make the separation of concerns clearer, and makes it clear what any kind of "forecaster" must behave like (take as input, give as output)
There was a problem hiding this comment.
I would put these in .forecasters.base_module, .forecasters.deterministic_module and .forecasters.probablistic_module
There was a problem hiding this comment.
any shared functionality between the deterministic and probabilistic forecaster modules would then sit in the base forecaster module
There was a problem hiding this comment.
I am wondering whether "Forecaster" is actually a good term as this limits the usage to forecasts, but we might also have negative timestepping at some point or do downscaling. So is "Predictor" maybe a better/more general term here?
There was a problem hiding this comment.
Restructured the module hierarchy per @leifdenby's suggestion:
-
Introduced
BaseForecasterModule(pl.LightningModule, ABC)inneural_lam/models/forecasters/base_module.py, holding everything shared between evaluation modes:training_step,common_step,on_after_batch_transfer(batch standardization),_warn_skipped_val_steps,on_validation_epoch_end,plot_examples,aggregate_and_plot_metrics,on_load_checkpoint.validation_step,test_step, andon_test_epoch_endare left@abstractmethodsince they differ meaningfully between deterministic and ensemble evaluation. -
Renamed
ForecasterModule→DeterministicForecasterModuleand moved it toneural_lam/models/forecasters/deterministic_module.py. -
Moved
ProbabilisticForecasterModuletoneural_lam/models/forecasters/probabilistic_module.py. It now implementsBaseForecasterModuledirectly as a sibling of the deterministic module, rather than subclassing it. This also resolves the tension from the earliersuper().validation_step()discussion since both modules independently implement the same abstract contract, there's no more inheritance-reuse temptation (and no risk of wasting a forward pass or reintroducing the loss/diagnostic-naming conflation that came up earlier in review).
All 256 tests and pre-commit hooks (black/isort/flake8/interrogate/mypy) pass. Touched 16 files total (mostly mechanical import/rename updates in tests and docstrings).
Commit: d020691.
There was a problem hiding this comment.
Good point, worth thinking through. Two observations that make me lean toward keeping "Forecaster" for now, though I'm open to it:
-
There's already a
StepPredictorin this same package the single-step model thatARForecasterwraps and unrolls into a multi-step trajectory. Right now "Forecaster wraps StepPredictor" reads unambiguously as two different levels of abstraction. IfForecasterbecomesPredictor, we'd havePredictor(multi-step) sitting next toStepPredictor(single-step), the containment relationship gets muddier, since the names no longer signal which one is "bigger." -
Backward timestepping and downscaling are worth keeping the door open for, but I think even those still fit "produce a forecast/estimate over a target domain from some initial condition" reasonably well. "Predictor" being more generic also makes it vaguer about what is being predicted, which the current interface (states at future grid times) is actually fairly specific about.
So unless we're also willing to rename StepPredictor at the same time to keep the two levels distinct (e.g. Forecaster/StepModel), I'd rather not do a naming-only rename across the whole package without a concrete near-term use case (downscaling, hindcasting) driving it. Happy to revisit this once one of those lands.
There was a problem hiding this comment.
This was a very good suggestion, fully agree. Also good to keep the classes in separate files.
This also resolves the tension from the earlier super().validation_step() discussion since both modules independently implement the same abstract contract
Indeed, that's very nice!
There was a problem hiding this comment.
As discussed a bit on slack, it seems best to wait with any renaming of these classes and stick to Forecaster now.
There was a problem hiding this comment.
I also think the files containing the ForecasterModules should be in their own modules directory, and not mixed with the forecasters.
|
Thank you for your review ! I am on it 🤩 |
Rename the ensemble-mean diagnostic keys from *_loss_unroll/*_mean_loss to *_ens_rmse_unroll/*_mean_ens_rmse so they aren't conflated with the training loss, per review feedback.
Remove explanatory comments around the per_var_std overrides; the assignments are clear on their own.
…crete deterministic/probabilistic modules Introduce BaseForecasterModule (abstract) under models/forecasters/ holding shared plumbing (training_step, common_step, batch standardization, checkpoint compatibility, plotting/aggregation helpers), with validation_step, test_step and on_test_epoch_end left abstract since they differ meaningfully between evaluation modes. Rename ForecasterModule to DeterministicForecasterModule and move it, alongside ProbabilisticForecasterModule, into forecasters/ as siblings implementing the shared contract, rather than one subclassing the other.
|
@leifdenby @joeloskarsson @observingClouds, I hope these commits solve most of the issue, where i need a little more direction, I have left the comments asking for more help ! Thank you hope to here from you all soon. 🤩 |
Removes estimate_likelihood/compute_step_loss and the per_var_std buffer they existed to feed, plus the now-unused config constructor arg (its only use was building per_var_std). Per the objective now living on the Forecaster (mllam#700), the predictor stays a pure network construct: encoder/prior/decoder plus the forward sampling path.
|
Read through all of the fixes and discussion here now. Some small follow-up comments on my earlier points, but in general I think this is looking quite good. The refactor into a base module class was very good, but also moved some things around a bit. I would be happy to give this another full read through before approving once I am back from vacation. I think that should be timely if this goes in in v0.8.0. |
Add JetBrains IDE project directory to the ignore list alongside the existing .vim/.vscode entries.
…aster DeterministicForecasterModule.validation_step/test_step read self.forecaster.loss and self.forecaster.per_var_std directly, so the Module still knew how to compute a loss from a prediction. Add an abstract Forecaster.score (implemented on ARForecaster) that resolves the pred_std fallback and applies a scoring rule internally; the Module now only calls forecaster.score(...) and never touches loss/per_var_std itself.
…ster BaseForecasterModule.__init__ raised if forecaster.per_var_std was None and the forecaster didn't predict its own std -- a check on the Forecaster's configuration living in the wrong class. Move it into ARForecaster via a new _resolve_pred_std helper, shared by score() and compute_training_loss(): construction with config=None now always succeeds (a valid state for forecasters only ever used for inference), and the ValueError instead fires from the Forecaster itself, only once scoring is attempted without any std to use.
…dules/ BaseForecasterModule, DeterministicForecasterModule and ProbabilisticForecasterModule (Lightning wrappers around a Forecaster) were mixed in with the Forecaster classes themselves under neural_lam/models/forecasters/. Move them into their own neural_lam/models/modules/ package (base.py, deterministic.py, probabilistic.py) so the two concerns: what a forecaster is vs how it's trained/evaluated by Lightning live in separate directories. Pure move: internal imports and neural_lam/models/__init__.py updated accordingly, no behavioural change.
Resolves the conflict on neural_lam/models/module.py: main's PR mllam#675 added --train_steps_to_log and deduplicated the val/test/train prediction+loss+logging pattern into _compute_prediction_and_loss / _log_step_loss, on top of the pre-mllam#700 module.py where the Module still owned self.loss/self.per_var_std directly. This branch had since deleted that file (split into modules/base.py, deterministic.py, probabilistic.py, with loss ownership moved onto the Forecaster). Ports mllam#675's dedup/logging helpers (_warn_skipped_steps, _log_step_loss, the train_steps_to_log hparam) into modules/base.py, and rewires DeterministicForecasterModule's validation_step/test_step to use them via forecaster.score() instead of the removed self.loss/self.per_var_std. _compute_prediction_and_loss is dropped rather than ported: it assumes a single per-step-decomposable loss, which conflicts with compute_training_loss's contract that a forecaster's training objective may not decompose per step (e.g. an ELBO). training_step therefore still only logs the aggregate train_loss; train_steps_to_log is accepted for CLI/checkpoint compatibility but does not yet produce a train_loss_unroll{i} breakdown -- documented inline as a deliberate scope boundary, not an oversight. Also fixes two call sites (deterministic.py, probabilistic.py) that referenced the pre-rename _warn_skipped_val_steps, and absorbs the unrelated utils.py -> utils/ package split from mllam#682 (auto-merged cleanly, this branch never touched utils.py).
Describe your changes
This PR implements the abstract constructions agreed in the #685 discussion (RFC: a general probabilistic forecasting interface), deliberately without any concrete Graph-EFM instantiation, so the design can be reviewed on its own and this can server as the base for addition of any probabilistic model, not just
graph_efm.Move ownership of the training objective from
ForecasterModuleonto theForecaster. A new abstractForecaster.compute_training_lossreturns a finished(loss, loss_components)pair, so each forecaster owns its complete training objective (including assembling it from any internal terms).ForecasterModule.training_stepnow only injects the configured scoring rule (--loss), the interior mask andper_var_std, then logs the returned loss and components (component names prefixed with the phase). Per the discussion,this move is applied to the deterministic setup as well, so the same concept sits in the same place in both stacks: the deterministic
ARForecastertraining loss is unchanged in value (covered by an equality test), it is just computed by the forecaster itself. Note for reviewers: since the method is abstract onForecaster, any out-of-treeForecastersubclass now has to implement it; all in-repo forecasters go throughARForecaster.Add the probabilistic side, built on top of the deterministic one (no sample dimension leaks into deterministic components):
ProbabilisticForecaster(abstract): declaressample_ensemble, producing members stacked along a new dimension after batch,(B, S, pred_steps, num_grid_nodes, d_state). This encodes the module's only assumption "the forecaster can create ensemble forecasts of the correct shape" as a type contract, and leaves room for non-AR implementations (diffusion/flow) later.ProbabilisticARForecaster(ARForecaster, ProbabilisticForecaster): unrolls independent trajectories through a step predictor that samples its output (sequential loop as the safe first default), and by default trains on the configured scoring rule applied to the ensemble mean. Model-specific objectives (e.g. Graph-EFM's ELBO + CRPS) will live in subclasses that overridecompute_training_loss.ProbabilisticForecasterModule: training inherited unchanged; validation samples aneval_ensemble_sizeensemble and logs the RMSE of the ensemble mean (the example ensemble metric suggested in RFC: a general probabilistic forecasting interface #685, standing in until the ensemble-metrics PR), reusing theval_mean_loss/val_loss_unroll{i}log keys so existing checkpoint callbacks work unchanged.test_stepraisesNotImplementedErrorrather than silently evaluating a single member deterministically; ensemble test evaluation and plotting are a follow-up.Dependencies: none. In particular this PR does not depend on the planned ensemble metrics (
crps_ens,spread_squared). Follow-ups: ensemble metrics, the Graph-EFM forecaster (ELBO + CRPS, its ownkl_beta/crps_weightconfig), ensemble test evaluation + plotting, andtrain_model.pywiring forprobabilistic models.
Issue Link
Graph-EFM instantiation and ensemble evaluation follow in later PRs).
prob_model_lamonmain, see issue Merge Graph-EFM model fromprob_model_lambranch #62Type of change
Checklist before requesting a review
pullwith--rebaseoption if possible).Checklist for reviewers
Each PR comes with its own improvements and flaws. The reviewer should check the following:
Author checklist after completed review
reflecting type of change (add section where missing):
Checklist for assignee